import os

from PIL import Image
from tqdm import tqdm

from option import parse_args
from abnormal.multiple_actors_statistics import cmp_video_name


def get_min_width(video_path):
    w_min, w_max = 300, 0
    for image_name in os.listdir(video_path):
        with Image.open(os.path.join(video_path, image_name)) as img:
            w_min = min(w_min, img.size[0])
            w_max = max(w_max, img.size[0])
    return w_min, w_max


def get_width_less_than_100(video_path):
    width = {}
    for image_name in os.listdir(video_path):
        with Image.open(os.path.join(video_path, image_name)) as img:
            if img.size[0] < 100:
                if img.size[0] not in width:
                    width[img.size[0]] = 1
                else:
                    width[img.size[0]] += 1
    return width


def abnormal_crops(crops, func, template, ):
    video_names = os.listdir(crops)
    video_names.sort(key=cmp_video_name)
    for video_name in tqdm(video_names):
        video_name_split = video_name.split('_')
        if len(video_name_split) != 3:
            continue

        video_path = os.path.join(crops, video_name)
        # 1. Count with min and max
        w_min, w_max = get_min_width(video_path)

        if w_min < 100 or w_max > 350:
            func(template.format(video_name, w_min, w_max))

        # 2. Count with dict
        # width = get_width_less_than_100(video_path)
        # if len(width) != 0:
        #     print('{},{}'.format(video_name, width))


def main():
    args = parse_args()
    crops = os.path.join(args.root_dir, 'crops')
    args.save = True
    if args.save:
        with open('abnormal_width_statistics.dat', 'w') as f:
            abnormal_crops(crops, f.write, '{},{},{}\n')
    else:
        abnormal_crops(crops, print, '{},{},{}')


if __name__ == '__main__':
    main()
